{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "from glob import glob\n",
    "from pathlib import Path\n",
    "\n",
    "import promptquality as pq\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from llm_handler import LLMHandler\n",
    "\n",
    "from IPython.display import clear_output\n",
    "\n",
    "load_dotenv(\"../.env\")\n",
    "pq.login('https://console.demo.rungalileo.io')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "project_name = \"agent-lb-v1\"\n",
    "datasets = [Path(file_path).stem for file_path in glob(\"../data/datasets/*.parquet\")]\n",
    "print(len(datasets))\n",
    "PROJECT_ID = pq.get_project_from_name(project_name).id\n",
    "llm_handler = LLMHandler()\n",
    "models = llm_handler.get_all_models()\n",
    "datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "runs = pq.get_project_from_name(project_name).runs\n",
    "\n",
    "run_dict = defaultdict(list)\n",
    "for run in runs:\n",
    "    model, dataset = run.name.split(\" \")\n",
    "    run_dict[model] += [dataset]\n",
    "    \n",
    "for key, value in run_dict.items():\n",
    "    print(key, len(value))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import lru_cache\n",
    "\n",
    "@lru_cache(maxsize=1000)\n",
    "def get_model_score_for_dataset(model, dataset):\n",
    "    print(f\"Getting metrics for {model} {project_name} for dataset {dataset}\")\n",
    "    run_name = f\"{model} {dataset}\"\n",
    "    run_id = pq.get_run_from_name(run_name, PROJECT_ID).id\n",
    "    rows = pq.get_rows(\n",
    "        project_id=PROJECT_ID,\n",
    "        run_id=run_id,\n",
    "        starting_token=0,\n",
    "        limit=1000,\n",
    "    )\n",
    "    \n",
    "    scores = [round(d.metrics.tool_selection_quality, 2) for d in rows if d.metrics.tool_selection_quality is not None]\n",
    "    rationales = [d.metrics.tool_selection_quality_rationale for d in rows if d.metrics.tool_selection_quality is not None]\n",
    "    explanations = [d.metrics.tool_selection_quality_explanation for d in rows if d.metrics.tool_selection_quality is not None]\n",
    "    mean_score = round(np.mean(scores), 2)\n",
    "    return {\n",
    "        \"mean_score\": mean_score,\n",
    "        \"scores\": scores,\n",
    "        \"rationales\": rationales,\n",
    "        \"explanations\": explanations,\n",
    "    }\n",
    "\n",
    "# data = get_model_score_for_dataset(models[0], datasets[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "import threading\n",
    "from queue import Queue\n",
    "from typing import List, Dict\n",
    "import time\n",
    "from concurrent.futures import ThreadPoolExecutor\n",
    "from dataclasses import dataclass\n",
    "from threading import Semaphore\n",
    "from tqdm.auto import tqdm\n",
    "from tqdm.contrib.concurrent import thread_map\n",
    "\n",
    "@dataclass\n",
    "class ScoringTask:\n",
    "    model: str\n",
    "    dataset: str\n",
    "\n",
    "class ModelScorer:\n",
    "    def __init__(self, models: List[str], datasets: List[str], max_concurrent_datasets: int = 3):\n",
    "        self.models = models\n",
    "        self.datasets = datasets\n",
    "        self.results: Dict[tuple, float] = {}\n",
    "        self.results_lock = threading.Lock()\n",
    "        self.dataset_semaphore = Semaphore(max_concurrent_datasets)\n",
    "        self.progress_bars = {}\n",
    "        \n",
    "    def create_model_progress_bar(self, model: str, total_datasets: int) -> tqdm:\n",
    "        \"\"\"Create a progress bar for a specific model\"\"\"\n",
    "        return tqdm(\n",
    "            total=total_datasets,\n",
    "            desc=f\"Model {model}\",\n",
    "            position=self.models.index(model),\n",
    "            leave=True\n",
    "        )\n",
    "        \n",
    "    def process_task(self, task: ScoringTask):\n",
    "        \"\"\"\n",
    "        Process a single scoring task with semaphore control and progress tracking.\n",
    "        \"\"\"\n",
    "        with self.dataset_semaphore:\n",
    "            # Get or create progress bar for this model\n",
    "            with self.results_lock:\n",
    "                if task.model not in self.progress_bars:\n",
    "                    self.progress_bars[task.model] = self.create_model_progress_bar(\n",
    "                        task.model, \n",
    "                        len(self.datasets)\n",
    "                    )\n",
    "            \n",
    "            # Process the task\n",
    "            score = get_model_score_for_dataset(task.model, task.dataset)[\"mean_score\"]\n",
    "            \n",
    "            # Update results and progress\n",
    "            with self.results_lock:\n",
    "                self.results[(task.model, task.dataset)] = score\n",
    "                self.progress_bars[task.model].update(1)\n",
    "                self.progress_bars[task.model].set_postfix(\n",
    "                    last_dataset=task.dataset,\n",
    "                    last_score=f\"{score:.3f}\"\n",
    "                )\n",
    "\n",
    "    def run_scoring(self) -> Dict[tuple, float]:\n",
    "        \"\"\"\n",
    "        Run the scoring process for all model-dataset combinations with progress tracking.\n",
    "        \"\"\"\n",
    "        tasks = [\n",
    "            ScoringTask(model, dataset)\n",
    "            for model in self.models\n",
    "            for dataset in self.datasets\n",
    "        ]\n",
    "        \n",
    "        # Create overall progress bar\n",
    "        total_tasks = len(tasks)\n",
    "        \n",
    "        # Clear any existing progress bars\n",
    "        for _ in range(len(self.models)):\n",
    "            print()\n",
    "            \n",
    "        # Use ThreadPoolExecutor to manage the thread pool\n",
    "        with ThreadPoolExecutor(max_workers=min(len(tasks), 10)) as executor:\n",
    "            # Submit all tasks\n",
    "            futures = [executor.submit(self.process_task, task) for task in tasks]\n",
    "            \n",
    "            # Create overall progress bar\n",
    "            with tqdm(\n",
    "                total=total_tasks,\n",
    "                desc=\"Overall Progress\",\n",
    "                position=len(self.models),\n",
    "                leave=True\n",
    "            ) as overall_pbar:\n",
    "                # Wait for all tasks to complete\n",
    "                for future in futures:\n",
    "                    future.result()\n",
    "                    overall_pbar.update(1)\n",
    "        \n",
    "        # Close all progress bars\n",
    "        for pbar in self.progress_bars.values():\n",
    "            pbar.close()\n",
    "                \n",
    "        return self.results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "models = [\"gemini-2.5-pro-preview-03-25\"]\n",
    "\n",
    "scorer = ModelScorer(models, datasets, max_concurrent_datasets=3)\n",
    "\n",
    "print(\"Starting scoring process...\")\n",
    "start_time = time.time()\n",
    "\n",
    "results = scorer.run_scoring()\n",
    "\n",
    "end_time = time.time()\n",
    "print(f\"\\nScoring completed in {end_time - start_time:.2f} seconds\")\n",
    "\n",
    "# Print results in a formatted way\n",
    "print(\"\\nFinal Scores:\")\n",
    "print(\"-\" * 40)\n",
    "for (model, dataset), score in sorted(results.items()):\n",
    "    print(f\"{model:10} | {dataset:10} | {score:.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the initial DataFrame\n",
    "results_matrix = {\n",
    "    dataset: {model: scorer.results.get((model, dataset), None) \n",
    "                for model in scorer.models}\n",
    "    for dataset in scorer.datasets\n",
    "}\n",
    "df = pd.DataFrame(results_matrix).round(3)\n",
    "\n",
    "# Create average of toolace_single_func_call datasets\n",
    "toolace_cols = ['toolace_single_func_call_1', 'toolace_single_func_call_2']\n",
    "if all(col in df.columns for col in toolace_cols):\n",
    "    # Calculate average of the two datasets\n",
    "    df['toolace_single_func_call'] = df[toolace_cols].mean(axis=1)\n",
    "    # Drop the original columns\n",
    "    df = df.drop(columns=toolace_cols)\n",
    "\n",
    "# Calculate averages\n",
    "df['Model Avg'] = df.mean(axis=1)  # Average for each model across datasets\n",
    "dataset_avgs = df.mean()  # Average for each dataset across models\n",
    "\n",
    "# Add dataset averages as a new row\n",
    "df.loc['Dataset Avg'] = dataset_avgs\n",
    "\n",
    "# Sort by model performance (descending)\n",
    "df_sorted = pd.concat([\n",
    "    df.iloc[:-1].sort_values('Model Avg', ascending=False),\n",
    "    df.iloc[-1:] # Add back the Dataset Avg row at the end\n",
    "])\n",
    "\n",
    "# Optionally round all values to 3 decimal places for clean display\n",
    "df_sorted = df_sorted.round(3)\n",
    "df_sorted = df_sorted[[\"Model Avg\"] + [col for col in df_sorted.columns if col != \"Model Avg\"]]\n",
    "df_sorted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langgraph",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
